# 导包
from tensorflow.keras.preprocessing.image import ImageDataGenerator

""""
图片生成器 tf.keras.preprocessing.image.ImageDataGenerator
用到的参数：
-rescale:输入一个参数，通常为1/255  让像素值归一化
-directory:文件路径
-target_size:图像宽高缩放到指定大小
-batch_size:每次读取的图片数（默认32）
-class_mode:类别格式，默认'categorical'
如果是‘sparse':类别['paper','rock','scissor'-->[0,1,2]
categorical:类别['paper','rock','scissor'-->[1,0,0],[0,1,0],[0,0,1]
input:类别['paper','rock','scissor']保持不变
如果是None:不返回标签
"""
# 验证
"""
def train_val_generator(data_dir, target_size, batch_size, class_mode=None, subset='training'):
      train_val_datagen = ImageDataGenerator(rescale=1./255., validation_split=0.2)
      return train_val_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode,
        subset=subset
    )
"""
def train_val_generator(data_dir, target_size, batch_size, class_mode=None, subset='training'):
      train_val_datagen = ImageDataGenerator(rescale=1./255.)
      return train_val_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode,
        subset=subset
    )

"""
def test_generator(data_dir,target_size,batch_size,class_mode=None):
    test_datagen = ImageDataGenerator(rescale=1./255.)
    return test_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode,
    )
"""


# 测试
def test_generator(data_dir,target_size,batch_size,class_mode=None):
    test_datagen = ImageDataGenerator(rescale=1./255.)
    return test_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode,
    )

# 预测
def pred_generator(data_dir,target_size,batch_size,class_mode=None):
    pred_datagen=ImageDataGenerator(rescale=1./255.)
    return pred_datagen.flow_from_directory(
        directory=data_dir,
        target_size=target_size,
        batch_size=batch_size,
        class_mode=class_mode,
    )